import os

import mmcv
# using a pre-trained detector. 预训练配置
from mmcv import Config


root_path=r"/home/deepin/Documents/openmmlab/mmdetection/"

cfg_path=root_path+'configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py'

cfg = Config.fromfile(cfg_path)

from mmdet.apis import set_random_seed, inference_detector, show_result_pyplot

# Modify dataset type and path
cfg.dataset_type = 'KittiTinyDataset' # 自定义的数据集的 类
cfg.data_root = root_path+'data/kitti_tiny/'

cfg.data.test.type = 'KittiTinyDataset'
cfg.data.test.data_root = root_path+'data/kitti_tiny/'
cfg.data.test.ann_file = 'train.txt'
cfg.data.test.img_prefix = 'training/image_2' # 图片 所在文件夹

cfg.data.train.type = 'KittiTinyDataset'
cfg.data.train.data_root = root_path+'data/kitti_tiny/'
cfg.data.train.ann_file = 'train.txt'
cfg.data.train.img_prefix = 'training/image_2'

cfg.data.val.type = 'KittiTinyDataset'
cfg.data.val.data_root = root_path+'data/kitti_tiny/'
cfg.data.val.ann_file = 'val.txt'
cfg.data.val.img_prefix = 'training/image_2'

cfg.device='cuda' # 'ConfigDict' object has no attribute 'device'

# modify num classes of the model in box head 类别数
cfg.model.roi_head.bbox_head.num_classes = 3
# If we need to finetune a model based on a pre-trained detector, we need to
# use load_from to set the path of checkpoints. 预训练模型

cfg.load_from = root_path+'checkpoints/faster_rcnn_r50_caffe_fpn_mstrain_3x_coco_20210526_095054-1f77628b.pth'

# Set up working dir to save files and logs. 保存的文件夹
cfg.work_dir = 'work_dir'

# The original learning rate (LR) is set for 8-GPU training.
# We divide it by 8 since we only use one GPU.
#原始学习率（LR）被设置用于8-GPU训练。
#我们将其除以8，因为我们只使用一个GPU。
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# Change the evaluation metric since we use customized dataset.
#更改评估指标，因为我们使用自定义数据集。
cfg.evaluation.metric = 'mAP'

# We can set the evaluation interval to reduce the evaluation times
#我们可以设置评估间隔以减少评估时间
cfg.evaluation.interval = 12

# We can set the checkpoint saving interval to reduce the storage cost
#我们可以设置检查点保存间隔以降低存储成本
cfg.checkpoint_config.interval = 12

# Set seed thus the results are more reproducible #播种，结果更具可重复性
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# We can also use tensorboard to log the training process #我们还可以使用tensorboard记录培训过程
cfg.log_config.hooks = [
    dict(type='TextLoggerHook'),
    dict(type='TensorboardLoggerHook')]


# We can initialize the logger for training and have a look
# at the final config used for training
#我们可以初始化记录器进行培训并查看
#在用于培训的最终配置中
print(f'Config:\n{cfg.pretty_text}')

#保存 cfg 配置：
# 把自定义的 config 在 自定义的working dir 下 保存一份------------推理和测试 可以用到的
cfg.dump(F'{cfg.work_dir}/KittiTinyDataset_cfgformat.py')



#训练----------------------------------------------------
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector


# Build dataset 构建 数据集
datasets = [build_dataset(cfg.data.train)]

# Build the detector 构建 识别
model = build_detector(cfg.model)

# Add an attribute for visualization convenience 添加属性以方便可视化，模型的类别
model.CLASSES = datasets[0].CLASSES

# Create work_dir
mmcv.mkdir_or_exist(os.path.abspath(cfg.work_dir))
train_detector(model, datasets, cfg, distributed=False, validate=True) # -------进行----训练-----

# # # 查看 评估结果 load tensorboard in colab
# # %load_ext tensorboard
# #
# # # see curves in tensorboard
# # tensorboard --logdir ./tutorial_exps  #--------运行这个看看 数的 曲线



